from preprocessing.constants import OULU_NPU, CASIA_FASD, REPLAY_ATTACK
from .casia_fasd_dataset import get_casia_fasd_dataloader, get_casia_fasd_test_dataloader
from .oulu_npu_dataset import get_oulu_npu_dataloader, get_oulu_npu_test_dataloader
from .oulu_npu_dataset import get_oulu_npu_dataloader_leave_one, get_oulu_npu_test_dataloader_leave_one
from .replay_attack_dataset import get_replay_attack_dataloader, get_replay_attack_test_dataloader


def get_dataloader(model_cfg, args):
    if args.prefix == OULU_NPU:
        if args.protocol == 'Protocol_3' or args.protocol == 'Protocol_4':
            train_sampler, train_loader, val_loader = \
                get_oulu_npu_dataloader_leave_one(model_cfg, args, color_space=args.color_space)
        else:
            train_sampler, train_loader, val_loader = \
                get_oulu_npu_dataloader(model_cfg, args, color_space=args.color_space)
    elif args.prefix == CASIA_FASD:
        train_sampler, train_loader, val_loader = \
            get_casia_fasd_dataloader(model_cfg, args, color_space=args.color_space)
    elif args.prefix == REPLAY_ATTACK:
        train_sampler, train_loader, val_loader = \
            get_replay_attack_dataloader(model_cfg, args, color_space=args.color_space)
    return train_sampler, train_loader, val_loader


def get_test_dataloader(model_cfg, args):
    if args.prefix == OULU_NPU:
        if args.protocol == 'Protocol_3' or args.protocol == 'Protocol_4':
            test_loader = get_oulu_npu_test_dataloader_leave_one(model_cfg, args, color_space=args.color_space)
        else:
            test_loader = get_oulu_npu_test_dataloader(model_cfg, args, color_space=args.color_space)
    elif args.prefix == CASIA_FASD:
        test_loader = get_casia_fasd_test_dataloader(model_cfg, args, color_space=args.color_space)
    elif args.prefix == REPLAY_ATTACK:
        test_loader = get_replay_attack_test_dataloader(model_cfg, args, color_space=args.color_space)
    return test_loader
